from colorama import Fore
from camel.societies import RolePlaying
from camel.utils import print_text_animated
from dotenv import load_dotenv
from camel.agents import ChatAgent
from camel.models import ModelFactory
from camel.types import ModelPlatformType
from camel.messages import OpenAIMessage
from camel.logger import get_logger, set_log_file, set_log_level
import os

# Set log output to a file using an absolute path
log_path = os.path.join(
    os.path.dirname(os.path.abspath(__file__)),
    'role_playing.log',
)
print(log_path)
set_log_file(log_path)

# Set the logging level
set_log_level('INFO')

load_dotenv(override=True)

base_url = os.getenv("base_url")
api_key = os.getenv("api_key")

def main(model=None, chat_turn_limit=50) -> None:
    task_prompt = "详细设计大模型学习计划"
    role_play_session = RolePlaying(
        assistant_role_name="老师",
        assistant_agent_kwargs=dict(model=model),
        user_role_name="学生",
        user_agent_kwargs=dict(model=model),
        task_prompt=task_prompt,
        with_task_specify=True,
        task_specify_agent_kwargs=dict(model=model),
        output_language="中文"
    )

    print(
        Fore.GREEN
        + f"AI Assistant sys message:\n{role_play_session.assistant_sys_msg}\n"
    )
    print(
        Fore.BLUE + f"AI User sys message:\n{role_play_session.user_sys_msg}\n"
    )

    print(Fore.YELLOW + f"Original task prompt:\n{task_prompt}\n")
    print(
        Fore.CYAN
        + "Specified task prompt:"
        + f"\n{role_play_session.specified_task_prompt}\n"
    )
    print(Fore.RED + f"Final task prompt:\n{role_play_session.task_prompt}\n")

    n = 0
    input_msg = role_play_session.init_chat()
    while n < chat_turn_limit:
        n += 1
        assistant_response, user_response = role_play_session.step(input_msg)

        if assistant_response.terminated:
            print(
                Fore.GREEN
                + (
                    "AI Assistant terminated. Reason: "
                    f"{assistant_response.info['termination_reasons']}."
                )
            )
            break
        if user_response.terminated:
            print(
                Fore.GREEN
                + (
                    "AI User terminated. "
                    f"Reason: {user_response.info['termination_reasons']}."
                )
            )
            break

        print_text_animated(
            Fore.BLUE + f"AI User:\n\n{user_response.msg.content}\n"
        )
        print_text_animated(
            Fore.GREEN + "AI Assistant:\n\n"
            f"{assistant_response.msg.content}\n"
        )

        if "CAMEL_TASK_DONE" in user_response.msg.content:
            break

        input_msg = assistant_response.msg


if __name__ == "__main__":
    model = ModelFactory.create(
        model_platform=ModelPlatformType.OPENAI_COMPATIBLE_MODEL,
        model_type="DeepSeek-R1",
        url=base_url,
        api_key=api_key
    )
    main(model=model)